package io.gitee.zhangbinhub.admin.gateway.conf

import com.fasterxml.jackson.databind.ObjectMapper
import io.gitee.zhangbinhub.admin.common.constant.RouteConstant
import io.gitee.zhangbinhub.admin.gateway.constant.GateWayConstant
import io.gitee.zhangbinhub.admin.gateway.consumer.UpdateRouteConsumerProcess
import io.gitee.zhangbinhub.admin.gateway.log.RouteLogBridge
import io.gitee.zhangbinhub.admin.gateway.log.RouteLogService
import org.reactivestreams.Publisher
import org.springframework.beans.factory.annotation.Autowired
import org.springframework.boot.autoconfigure.AutoConfigureBefore
import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean
import org.springframework.cloud.gateway.filter.GatewayFilterChain
import org.springframework.cloud.gateway.filter.GlobalFilter
import org.springframework.cloud.gateway.filter.NettyWriteResponseFilter
import org.springframework.cloud.stream.config.BindingServiceConfiguration
import org.springframework.cloud.stream.function.StreamBridge
import org.springframework.context.annotation.Bean
import org.springframework.context.annotation.Configuration
import org.springframework.core.Ordered
import org.springframework.core.io.buffer.DataBuffer
import org.springframework.core.io.buffer.DataBufferFactory
import org.springframework.http.HttpHeaders
import org.springframework.http.MediaType
import org.springframework.http.server.reactive.ServerHttpResponseDecorator
import org.springframework.web.server.ServerWebExchange
import reactor.core.publisher.Flux
import reactor.core.publisher.Mono

/**
 * @since JDK 11
 */
@Configuration(proxyBeanMethods = false)
@AutoConfigureBefore(BindingServiceConfiguration::class)
class RouteConfiguration @Autowired constructor(private val routeLogService: RouteLogService) {
    @Bean(RouteConstant.UPDATE_ROUTE_CONSUMER)
    @ConditionalOnMissingBean(name = [RouteConstant.UPDATE_ROUTE_CONSUMER])
    fun updateRouteConsumer(updateRouteConsumerProcess: UpdateRouteConsumerProcess): (String) -> Unit {
        return { updateRouteConsumerProcess.accept(it) }
    }

    @Bean
    fun routeLogBridge(streamBridge: StreamBridge, objectMapper: ObjectMapper) =
        RouteLogBridge(streamBridge, objectMapper, RouteConstant.ROUTE_LOG_OUTPUT)

    @Bean
    fun logRequestFilter(routeLogBridge: RouteLogBridge): GlobalFilter = GlobalFilter { exchange, chain ->
        exchange.mutate().request(exchange.request.mutate().headers { httpHeaders: HttpHeaders ->
            httpHeaders.add(GateWayConstant.GATEWAY_HEADER_REQUEST_TIME, System.currentTimeMillis().toString())
        }.build()).build().let { serverWebExchange ->
            chain!!.filter(serverWebExchange.also {
                routeLogBridge.sendMessage(routeLogService.createRouteLogMessage(it))
            })
        }
    }

    @Bean
    fun logResponseFilter(routeLogBridge: RouteLogBridge): GlobalFilter = object : GlobalFilter, Ordered {
        override fun filter(exchange: ServerWebExchange?, chain: GatewayFilterChain?): Mono<Void> {
            val bufferFactory: DataBufferFactory = exchange!!.response.bufferFactory()
            val decoratedResponse: ServerHttpResponseDecorator =
                object : ServerHttpResponseDecorator(exchange.response) {
                    override fun writeWith(body: Publisher<out DataBuffer?>): Mono<Void?> {
                        val routeLogMessage = routeLogService.createRouteLogMessage(exchange)
                        val responseTime = System.currentTimeMillis()
                        routeLogMessage.processTime = responseTime - routeLogMessage.requestTime!!
                        routeLogMessage.responseTime = responseTime
                        exchange.response.statusCode?.let {
                            routeLogMessage.responseStatus = it.value()
                        }
                        return if (routeLogMessage.applyToken && body is Flux<*> &&
                            (exchange.response.headers.getFirst(HttpHeaders.CONTENT_TYPE)
                                ?: "").contains(MediaType.APPLICATION_JSON_VALUE, true)
                        ) {
                            val fluxBody = body as Flux<out DataBuffer?>
                            super.writeWith(fluxBody.buffer().map { dataBuffers ->
                                routeLogService.afterRoute(routeLogMessage, dataBuffers).let {
                                    routeLogBridge.sendMessage(routeLogMessage)
                                    bufferFactory.wrap(it)
                                }
                            })
                        } else {
                            routeLogService.afterRoute(routeLogMessage).let {
                                routeLogBridge.sendMessage(routeLogMessage)
                                super.writeWith(body)
                            }
                        }
                    }
                }
            return chain!!.filter(exchange.mutate().response(decoratedResponse).build())
        }

        override fun getOrder(): Int {
            return NettyWriteResponseFilter.WRITE_RESPONSE_FILTER_ORDER - 1
        }
    }

}
